import torch
import torch.nn as nn
from torchvision import transforms
from PIL import Image
import matplotlib
import matplotlib.pyplot as plt

matplotlib.use('TkAgg')

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.linear1 = nn.Linear(784, 512)
        self.linear2 = nn.Linear(512, 256)
        self.linear3 = nn.Linear(256, 128)
        self.linear4 = nn.Linear(128, 64)
        self.linear5 = nn.Linear(64, 10)
        self.act = nn.ReLU()

    def forward(self, x):
        x = x.view(-1, 784)
        x = self.act(self.linear1(x))
        x = self.act(self.linear2(x))
        x = self.act(self.linear3(x))
        x = self.act(self.linear4(x))
        x = self.linear5(x)
        return x


model = Net().to(device)
model_path = "../models/mnist_model.pth"
model.load_state_dict(torch.load(model_path))
model.eval()

transform = transforms.Compose([
    transforms.Resize((28, 28)),
    transforms.Grayscale(),
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])


def predict_image(image_path, model, transform):
    image = Image.open(image_path)
    plt.figure(figsize=(6, 6))
    plt.imshow(image, cmap="gray")
    plt.title("Original Image")
    plt.axis('off')
    plt.show()
    image_tensor = transform(image).unsqueeze(0).to(device)
    with torch.no_grad():
        outputs = model(image_tensor)
        probabilities = torch.softmax(outputs, dim=1)
        predicted_class = torch.argmax(probabilities, dim=1).item()
        confidence = probabilities[0][predicted_class].item()
    print("\n所有类别的预测概率:")
    for i, prob in enumerate(probabilities[0]):
        print(f"数字 {i}: {prob:.4f} ({prob * 100:.2f}%)")
    print(f"\n预测结果: 数字 {predicted_class}")
    print(f"置信度: {confidence:.4f} ({confidence * 100:.2f}%)")
    return predicted_class, confidence


image_path = "../data/Mnist/test/8.png"
predicted_digit, confidence = predict_image(image_path, model, transform)
print(f"预测的数字是: {predicted_digit}")
print(f"置信度是: {confidence:.4f}")
